package jk.framework.spring.cloud.canary.loadbalancer.comonpent;



import jk.framework.canary.core.CanaryMgt;
import jk.framework.canary.core.po.RouteContext;
import jk.framework.canary.core.po.CanaryResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * @author cuichao
 * 金丝雀RoundRobin负载均衡器
 */
public class CanaryRoundRobinLoadBalancer extends RoundRobinLoadBalancer {

    private static final Logger log = LoggerFactory.getLogger(CanaryRoundRobinLoadBalancer.class);

    final AtomicInteger position;

    final String serviceId;

    final LoadBalanceRouteContextFactory contextFactory;

    ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;



    public CanaryRoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                        String serviceId,
                                        LoadBalanceRouteContextFactory contextFactory
                                        ) {
        super(serviceInstanceListSupplierProvider, serviceId);
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.contextFactory = contextFactory;
        this.position = new AtomicInteger(1000);
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map(serviceInstances -> processInstanceResponse(request,supplier, serviceInstances));
    }



    private Response<ServiceInstance> processInstanceResponse(Request request,ServiceInstanceListSupplier supplier, List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(request,serviceInstances);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(Request request ,List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            if (log.isDebugEnabled()) {
                log.debug("No servers available for service: " + serviceId);
            }
            return new EmptyResponse();
        }



        List<ServiceInstance> candidateInStances = new ArrayList<>();
        List<ServiceInstance> unCandidateInStances = new ArrayList<>();

        AtomicBoolean instanceHasAllow = new AtomicBoolean(false);

        instances.forEach(v -> {
            RouteContext<RequestData> context = contextFactory.getContext(v,request);
            CanaryResponse<ServiceInstance> matchResponse = CanaryMgt.process(context, v);
            if(matchResponse.isMatchVersion()){
                candidateInStances.add(v);
            }else {
                unCandidateInStances.add(v);
            }
            if(matchResponse.isAllow()){
                instanceHasAllow.set(true);
            }
        });

        List<ServiceInstance> matchInstances;
        if(instanceHasAllow.get()){
            matchInstances = candidateInStances;
        }else {
            matchInstances = unCandidateInStances;
        }

        if(log.isDebugEnabled()){
            log.debug(String.format("当前请求是否允许路由到灰度发布节点: %s",instanceHasAllow.get()));
            List<String> collect = matchInstances.stream().map(v -> v.getHost() + ":" + v.getPort()).collect(Collectors.toList());
            log.debug(String.format("符合路由的服务节点为:%s",collect));
        }

        // TODO: enforce order?
        int pos = Math.abs(this.position.incrementAndGet());

        ServiceInstance instance = matchInstances.get(pos % matchInstances.size());

        return new DefaultResponse(instance);
    }
}
